package cn.xxm.dongbao.user.secure.dos;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;

/**
 * Get a token from the TokenBucket to control the request frequency
 */
public class TokenBucket implements Serializable {
    private static final long serialVersionUID = 1567868875894230569L;
    private static final Logger LOG = LoggerFactory.getLogger(TokenBucket.class);
    private final Object lockObj = new Object();
    private int count;
    private Long lastUpdateTime;
    private int maxTokenNum = 10;
    private int period = 10000;

    public TokenBucket(DosConfig conf) {
        setMaxTokenNum(conf.getHits());
        setPeriod(conf.getPeriod());
        count = maxTokenNum;
        lastUpdateTime = System.currentTimeMillis();
    }

    /**
     * Validate TokenBucket has token or not.     *     * @return validate result
     */
    public boolean hasToken() {
        synchronized (lockObj) {
            long interval = System.currentTimeMillis() - lastUpdateTime;
            lastUpdateTime = System.currentTimeMillis();
            if (interval > 0) {
                long addTokenNum = interval * maxTokenNum / period;
                if ((count + addTokenNum) > maxTokenNum) {
                    count = maxTokenNum;
                } else {
                    count += (int) addTokenNum;
                }
            }
            // get token from bucket
            if (count > 0) {
                count--;
                return true;
            }
            return false;
        }
    }

    private void setMaxTokenNum(String hits) {
        try {
            int result = Integer.parseInt(hits);
            if (result > 0) {
                maxTokenNum = result;
            } else {
                LOG.info("The parameter hits config error, Use the default value: {}.", maxTokenNum);
            }
        } catch (NumberFormatException ex) {
            LOG.info("The parameter hits config error, Use the default value: {}.", maxTokenNum);
        }
    }

    private void setPeriod(String period) {
        try {
            int result = Integer.parseInt(period);
            if (result > 0) {
                this.period = result * 1000;
            } else {
                LOG.info("The parameter period config error, Use the default value: {}.", this.period);
            }
        } catch (NumberFormatException ex) {
            LOG.info("The parameter period config error, Use the default value: {}.", this.period);
        }
    }
}